{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import pickle\n",
    "import lasagne\n",
    "from lasagne.layers import *\n",
    "from lasagne.nonlinearities import *\n",
    "from lasagne.objectives import *\n",
    "from lasagne.init import HeNormal\n",
    "import theano\n",
    "from theano import tensor as T\n",
    "import sys\n",
    "import numpy as np\n",
    "from copy import deepcopy\n",
    "#enable importing of notebooks\n",
    "from nbfinder import NotebookFinder\n",
    "sys.meta_path.append(NotebookFinder())\n",
    "import inspect\n",
    "import os\n",
    "\n",
    "if \"THEANO_FLAGS\" in os.environ and \"gpu\" in os.environ[\"THEANO_FLAGS\"]:\n",
    "        from lasagne.layers import dnn\n",
    "from helper_fxns import get_detec_loss, softmax3D, softmax4D, AccuracyGetter\n",
    "import copy\n",
    "#if __name__ == \"__main__\":\n",
    "    #from data_loader import load_classification_dataset, load_detection_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def get_hyperparams(frame):\n",
    "    args, _, _, values = inspect.getargvalues(frame)\n",
    "    #return dict(zip(args,values))\n",
    "    #del values['frame']\n",
    "    return values\n",
    "\n",
    "def build_network(kwargs):\n",
    "    \n",
    "    '''Takes a pretrained classification net and adds a few convolutional layers on top of it\n",
    "    and defines a detection loss function'''\n",
    "    '''Args:\n",
    "                      \n",
    "                      num_convpool: number of conv layer-pool layer pairs\n",
    "                      delta: smoothing constant to loss function (ie sqrt(x + delta)) \n",
    "                            -> if x is 0 gradient is undefined\n",
    "                      num_filters\n",
    "                      num_fc_units\n",
    "                      num_extra_conv: conv layers to add on to each conv layer before max pooling\n",
    "                      nonlinearity: which nonlinearity to use throughout\n",
    "                      n_boxes: how many boxes should be predicted at each grid point,\n",
    "                      nclass: how many classes are we predicting,\n",
    "                      grid_size: size of the grid that encodes various \n",
    "                                locations of image (ie in the YOLO paper they use 7x7 grid)\n",
    "                      w_init: weight intitialization\n",
    "                      dropout_p: prob of dropping unit\n",
    "                      coord_penalty : penalty in YOLO loss function for getting coordinates wrong\n",
    "                      nonobj_penalty: penalty in YOLO loss for guessing object when there isn't one\n",
    "                      learning_rate\n",
    "                      weight_decay\n",
    "                      momentum\n",
    "                      load: whether to load weights or not\n",
    "                      load_path: path for loading weights'''\n",
    "\n",
    "    if kwargs[\"im_dim\"] == 3:\n",
    "        input_var = T.tensor5('input_var') \n",
    "    else:\n",
    "        input_var = T.tensor4('input_var')\n",
    "    target_var = T.tensor4('target_var') #is of shape (grid_size, grid_size,(n_boxes* 5 + nclass)\n",
    "    \n",
    "    print \"Building model and compiling functions...\" \n",
    "    \n",
    "    #make layers\n",
    "    networks = build_layers(input_var,kwargs)\n",
    "    \n",
    "    #load in any pretrained weights\n",
    "    if kwargs['ae_load_path'] != \"None\":\n",
    "        networks['ae'] = load_weights(kwargs['ae_load_path'], networks['ae'])\n",
    "        \n",
    "    if kwargs['yolo_load_path'] != \"None\":\n",
    "        networks['yolo'] = load_weights(kwargs['yolo_load_path'], networks['yolo'])\n",
    "    \n",
    "    #compile theano functions\n",
    "    fns = make_fns(networks, input_var, target_var, kwargs)\n",
    "     \n",
    "\n",
    "    return fns, networks\n",
    "\n",
    "\n",
    "def build_layers(input_var, nk):\n",
    "    '''nk: network_kwargs'''\n",
    "    '''conv, extra_convs, pool multiple times then fc with dropout, fc with dropout and softmax then reshape'''\n",
    "    \n",
    "    '''total number of conv layers is num_convpool * (1 + num_extra_conv)'''\n",
    "    \n",
    "    filter_dim = nk['filter_dim']\n",
    "    num_layers = nk['num_layers']\n",
    "    \n",
    "\n",
    "    filters_list = [128, 256, 512, 768, 1024, 1280]\n",
    "    conv = lasagne.layers.InputLayer(shape=nk['input_shape'])\n",
    "    for i in range(num_layers):\n",
    "\n",
    "        \n",
    "        num_filters = int(nk[\"filters_scale\"] * filters_list[i])\n",
    "\n",
    "        \n",
    "        if nk[\"im_dim\"] == 3:\n",
    "            conv = dnn.Conv3DDNNLayer(conv, num_filters=num_filters, filter_size=(3,5,5),pad=(1,2,2), stride=(1,2,2)) \n",
    "        else:\n",
    "            #print nk[\"filters_scale\"]\n",
    "            #print num_filters\n",
    "            conv = Conv2DLayer(conv, \n",
    "                                  num_filters=num_filters, \n",
    "                                  filter_size=nk['filter_dim'], \n",
    "                                  pad=nk['filter_dim'] / 2, \n",
    "                                  stride=2,\n",
    "                                  W=HeNormal(),\n",
    "                                  nonlinearity=LeakyRectify(0.1))\n",
    "\n",
    "        \n",
    "\n",
    "    encoder = conv\n",
    "    hid_fmap = conv\n",
    "    \n",
    " \n",
    "    if nk[\"yolo_batch_norm\"]:\n",
    "        encoder = BatchNormLayer(encoder)\n",
    "    \n",
    "    if nk[\"im_dim\"] == 3:\n",
    "            #encoder = FeaturePoolLayer(encoder, pool_size=640, axis=1)\n",
    "            box_conf = dnn.Conv3DDNNLayer(encoder, num_filters=2, filter_size=(3,3,3), \n",
    "                                          stride=(2,1,1), pad=(1,1,1), nonlinearity=softmax4D)\n",
    "    \n",
    "            #box_conf = NonlinearityLayer(box_conf, softmax4D)\n",
    "            class_conf = dnn.Conv3DDNNLayer(encoder, num_filters=nk['num_classes'], filter_size=(3,3,3), stride=(2,1,1), pad=(1,1,1), nonlinearity=softmax4D)\n",
    "            #class_conf = NonlinearityLayer(class_conf, softmax4D)\n",
    "            coord_net = dnn.Conv3DDNNLayer(encoder, num_filters=4, filter_size=(3,3,3), stride=(2,1,1), pad=(1,1,1), W=nk['w_init'],\n",
    "                            nonlinearity=rectify)\n",
    "            #outputs a batch_size x 10 x 4 x 12 x 18 \n",
    "            bbox_reg = ConcatLayer([coord_net,box_conf, class_conf])\n",
    "            #print get_output_shape(bbox_reg,input_shapes=(1,16,8,768,1152))\n",
    "            s = get_output_shape(bbox_reg)\n",
    "            # reshape to be like the 2D case -> batch_size*time_steps(4) x 10 x 12 x 18\n",
    "            net = ExpressionLayer(bbox_reg, function=lambda g: g.transpose((0,2,1,3,4)),\n",
    "                                  output_shape=(nk[\"batch_size\"],s[2], s[1],s[3], s[4]))\n",
    "            #print get_output_shape(net,input_shapes=(1,16,8,768,1152))\n",
    "            # after transpose\n",
    "            s = get_output_shape(net)\n",
    "            bbox_reg = ReshapeLayer(net, shape=(nk[\"batch_size\"],s[1], s[2],s[3], s[4]))\n",
    "            #print get_output_shape(bbox_reg,input_shapes=(1,16,8,768,1152))\n",
    "            \n",
    "            \n",
    "    \n",
    "    else:   \n",
    "        box_conf = Conv2DLayer(encoder, num_filters=2, filter_size=3, pad=1, nonlinearity=linear)\n",
    "        box_conf = NonlinearityLayer(box_conf, softmax3D)\n",
    "        class_conf = Conv2DLayer(encoder, num_filters=nk['num_classes'], filter_size=3, pad=1, nonlinearity=linear)\n",
    "        class_conf = NonlinearityLayer(class_conf, softmax3D)\n",
    "        coord_net = Conv2DLayer(encoder, num_filters=4, filter_size=3,pad=1, W=nk['w_init'],\n",
    "                                nonlinearity=rectify)\n",
    "    \n",
    "    \n",
    "        bbox_reg = ConcatLayer([coord_net,box_conf, class_conf])\n",
    "    \n",
    "    for layer in get_all_layers(conv)[::-1]:\n",
    "        \n",
    "        if nk[\"batch_norm\"]:\n",
    "            conv = batch_norm(conv)\n",
    "        if isinstance(layer, InputLayer):\n",
    "            break\n",
    "        \n",
    "        conv = InverseLayer(conv, layer)\n",
    "        \n",
    "\n",
    "    \n",
    "    \n",
    "    return {'yolo': bbox_reg, 'ae':conv, 'hid_fmap':hid_fmap}#, \"decoder\":decoder_layers}\n",
    "        \n",
    "\n",
    "def load_weights(pickle_file_path, network):\n",
    "    '''grabs weights from an npz file'''\n",
    "    old_params = pickle.load(open(pickle_file_path, 'r'))\n",
    "\n",
    "    set_all_param_values(network, old_params)\n",
    "    return network\n",
    "    \n",
    "\n",
    "def make_fns(networks,input_var, target_var, kwargs ):\n",
    "    '''Compiles theano train, test, box_fns'''\n",
    "    #deterministic determines whether to use dropout or not in forward pass\n",
    "    #transpose output to match what loss expects\n",
    "    for k,v in networks.iteritems():\n",
    "        kwargs['logger'].info(\"\\n\" + k + \": \\n\")\n",
    "        for lay in get_all_layers(v):\n",
    "            kwargs[\"logger\"].info(str(lay) + \", \" + str(get_output_shape(lay)))\n",
    "    \n",
    "    yolo = networks['yolo']\n",
    "    ae = networks['ae']\n",
    "    hid_fmap = networks['hid_fmap']\n",
    "    yolo_test_prediction = lasagne.layers.get_output(yolo, deterministic=True, inputs=input_var)\n",
    "    yolo_prediction = lasagne.layers.get_output(yolo, deterministic=False,inputs=input_var)\n",
    "    \n",
    "    ae_test_prediction = lasagne.layers.get_output(ae, deterministic=True,inputs=input_var)\n",
    "    ae_prediction = lasagne.layers.get_output(ae, deterministic=False,inputs=input_var)\n",
    "    hid_fmap_pred = get_output(hid_fmap, deterministic=True, inputs=input_var)\n",
    "    \n",
    "    def make_hid_fn():\n",
    "        hid_fn = theano.function([input_var], hid_fmap_pred)\n",
    "        return hid_fn\n",
    "    \n",
    "    def make_loss(yolo_pred, ae_pred):\n",
    "        yolo_loss, terms = make_yolo_loss(yolo_pred)\n",
    "        ae_loss = make_ae_loss(ae_pred)\n",
    "        \n",
    "        #just to make sure we don't compute this if we don't want to\n",
    "        if kwargs['lambda_ae'] == 0.:\n",
    "            loss = yolo_loss\n",
    "        else:\n",
    "            loss = yolo_loss + kwargs['lambda_ae'] * ae_loss\n",
    "        return loss, yolo_loss, ae_loss, terms\n",
    "    \n",
    "    def make_yolo_loss(pred):\n",
    "        yolo_loss, terms = get_detec_loss(pred, target_var, kwargs)\n",
    "        \n",
    "        weightsl2 = lasagne.regularization.regularize_network_params(yolo, lasagne.regularization.l2)\n",
    "        weight_decay_term = kwargs['weight_decay'] * weightsl2\n",
    "        loss = yolo_loss + weight_decay_term\n",
    "        return loss, terms\n",
    "    \n",
    "    def make_ae_loss(pred):\n",
    "        loss = squared_error(pred, input_var)\n",
    "        weightsl2 = lasagne.regularization.regularize_network_params(ae, lasagne.regularization.l2)\n",
    "        loss += kwargs['weight_decay'] * weightsl2\n",
    "        return loss.mean()\n",
    "        \n",
    "    \n",
    "    \n",
    "\n",
    "    def make_train_fn(test=False):\n",
    "        '''takes as input the input, target vars and ouputs a loss'''\n",
    "        if test:\n",
    "            loss, yolo_loss, ae_loss, terms =  make_loss(yolo_test_prediction, ae_test_prediction)\n",
    "        else:   \n",
    "            loss, yolo_loss, ae_loss, terms =  make_loss(yolo_prediction, ae_prediction)\n",
    "        \n",
    "        #coord_term, size_term, conf_term, no_obj_conf_term, xentropy_term = terms\n",
    "        terms = T.stack(terms)\n",
    "        #only using params from yolo here -> because decoder has no new params -> tied weights\n",
    "        if not test:\n",
    "            params = lasagne.layers.get_all_params(yolo, trainable=True) #+ lasagne.layers.get_all_params(ae, trainable=True)\n",
    "            updates = lasagne.updates.adam(loss, params,learning_rate=kwargs['learning_rate'])\n",
    "\n",
    "            fn = theano.function([input_var, target_var], [loss,ae_loss,\n",
    "                                                                yolo_loss, terms ], \n",
    "                                                                updates=updates)\n",
    "        else:\n",
    "            fn = theano.function([input_var, target_var], [loss,ae_loss,\n",
    "                                                    yolo_loss, terms ])\n",
    "        \n",
    "        def fxn(inp, target):\n",
    "            loss, ae_loss, yolo_loss, terms = fn(inp, target)\n",
    "            coord_term, size_term, conf_term, no_obj_conf_term, xentropy_term = terms\n",
    "            \n",
    "            return dict(loss=loss, ae_loss=ae_loss, \n",
    "                        yolo_loss=yolo_loss) \n",
    "                        \n",
    "                        \n",
    "                        #coord_term=coord_term, size_term=size_term, conf_term=conf_term, \n",
    "                        #no_obj_conf_term=no_obj_conf_term, xentropy_term=xentropy_term)\n",
    "            \n",
    "            \n",
    "            \n",
    "            \n",
    "        return fxn\n",
    "        \n",
    "    \n",
    "    def make_test_or_val_fn():\n",
    "        \n",
    "        return make_train_fn(test=True)\n",
    "    \n",
    "    \n",
    "\n",
    "    \n",
    "    def make_yolo_pred_fn():\n",
    "        '''takes as input the input, target vars and outputs the predicted grid'''\n",
    "        pred_fn = theano.function([input_var], yolo_test_prediction)\n",
    "        return pred_fn\n",
    "    \n",
    "    def make_ae_pred_fn():\n",
    "        pred_fn = theano.function([input_var], ae_test_prediction)\n",
    "        return pred_fn\n",
    "        \n",
    "    def make_box_fn():\n",
    "        pred_fn = make_yolo_pred_fn()\n",
    "        def box_fn(x, y, num_classes=kwargs['num_classes']):\n",
    "            y_tensor = y\n",
    "            pred_tensor = pred_fn(x)\n",
    "            acc_obj = AccuracyGetter(kwargs)\n",
    "            pred_boxes, gt_boxes = acc_obj.get_all_boxes(pred_tensor=pred_tensor, \n",
    "                                                 y_tensor=y_tensor)\n",
    "            return pred_boxes, gt_boxes\n",
    "        return box_fn\n",
    "    \n",
    "   \n",
    "        \n",
    "            \n",
    "    def make_acc_fn():\n",
    "        '''takes as input the input, target vars and outputs the predicted and the ground truth boxes)'''\n",
    "        pred_fn = make_yolo_pred_fn()\n",
    "        def acc_fn(inp, gt, conf_thresh=None, iou_thresh=None):\n",
    "            pred = pred_fn(inp)\n",
    "            if conf_thresh is not None:\n",
    "                kwargs[\"conf_thresh\"] = conf_thresh\n",
    "            if iou_thresh is not None:\n",
    "                kwargs[\"iou_thresh\"] = iou_thresh\n",
    "            acc_getter = AccuracyGetter(kwargs)\n",
    "            acc_dict = acc_getter.get_scores(pred,gt)\n",
    "            return acc_dict\n",
    "    \n",
    "        return acc_fn\n",
    "    \n",
    "    train_fn = make_train_fn()\n",
    "    test_or_val_fn = make_test_or_val_fn()\n",
    "    acc_fn = make_acc_fn()\n",
    "    yolo_pred_fn = make_yolo_pred_fn()\n",
    "    ae_pred_fn = make_ae_pred_fn()\n",
    "    box_fn = make_box_fn()\n",
    "    hid_fn = make_hid_fn()\n",
    "    return {\"tr\":train_fn, \"val\":test_or_val_fn, \"test\":test_or_val_fn, \"acc\": acc_fn, \"rec\": ae_pred_fn, \"box\": box_fn, \"hid\": hid_fn}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# import theano\n",
    "\n",
    "# a = theano.tensor.scalar()\n",
    "# b = theano.tensor.scalar()\n",
    "# c = theano.tensor.scalar()\n",
    "# x = theano.tensor.stack([a, b, c])\n",
    "\n",
    "# x.eval({a:2,b:3,c:4})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
